# 人工智能NLP-Agent数字人项目-04-基金数据问答任务工单V1.1-2.13
import csv
import sqlite3
import pandas as pd
import utils.configFinRAG as configFinRAG

def query_db(sql, cursor):
    """
    执行SQL查询并返回执行结果和成功标志
    :param sql: 要执行的SQL语句
    :param cursor: 数据库游标
    :return: 成功标志（1表示成功，0表示失败），执行结果
    """
    try:
        cursor.execute(sql)
        result = cursor.fetchall()
        # 将查询结果转换为字符串
        result_str = str(result)
        success_flag = 1
    except Exception as e:
        # 捕获异常并打印错误信息
        print(f"执行SQL语句时出错: {sql}\n错误信息: {e}")
        result_str = "error"
        success_flag = 0
    return success_flag, result_str

def main():
    # 读取生成的SQL文件
    try:
        question_sql_file = pd.read_csv(configFinRAG.question_sql_path, delimiter=",", header=0)
    except FileNotFoundError:
        print(f"未找到SQL文件: {configFinRAG.question_sql_path}")
        return
    except Exception as e:
        print(f"读取SQL文件时出错: {e}")
        return

    # 数据库文件路径
    db_path = r'C:\Users\langzi\Desktop\practical_training\project\local_number_human\Fay-fay-agent-edition0830\agent\tools\data\博金杯比赛数据.db'

    # 使用上下文管理器管理数据库连接和文件操作
    with sqlite3.connect(db_path) as conn, open(configFinRAG.question_sql_check_path, 'w', newline='', encoding='utf-8-sig') as file:
        cursor = conn.cursor()
        csvwriter = csv.writer(file)
        # 写入CSV文件的表头
        csvwriter.writerow(['问题id', '问题', 'SQL', 'flag', '执行结果'])

        # 遍历SQL文件中的每一行
        for index, row in question_sql_file.iterrows():
            sql = row['SQL']
            if sql != 'error':
                success_flag, exc_result = query_db(sql, cursor)
                csvwriter.writerow([row['问题id'], row['问题'], sql, success_flag, exc_result])

if __name__ == '__main__':
    main()